package com.chj.zhongji.class09.test01;

import java.util.ArrayList;
import java.util.List;
import java.util.PriorityQueue;

public class KthSmallest {
	public static int kthSmallest(int[][] mat, int k) {
		int m = mat.length;
//		int n = mat[0].length;
		int[] res = mat[0];
		for (int i = 1; i < m; i++) {
			res = helper(res, mat[i]);
		}
		return res[k - 1];
	}

	private static int[] helper(int[] nums1, int[] nums2) {
		List<Integer> temp = new ArrayList<>();
		int k = 200;
		PriorityQueue<int[]> pq = new PriorityQueue<>((o1, o2) -> Integer.compare(o1[0] + o1[1], o2[0] + o2[1]));
		int len1 = nums1.length;
		int len2 = nums2.length;
		for (int i = 0; i < len1; i++)
			pq.add(new int[] { nums1[i], nums2[0], 0 });
		while (k-- > 0 && !pq.isEmpty()) {
			int[] curr = pq.poll();
			temp.add(curr[0] + curr[1]);
			if (curr[2] < len2 - 1)
				pq.add(new int[] { curr[0], nums2[curr[2] + 1], curr[2] + 1 });
		}
		int[] res = new int[temp.size()];
		for (int i = 0; i < temp.size(); i++)
			res[i] = temp.get(i);
		return res;
	}

	public static void main(String[] args) {
		int[][] matrix = new int[][] { { 1, 10, 10 }, { 1, 4, 5 }, { 2, 3, 6 } };
		System.out.println(kthSmallest(matrix, 7));
	}
}
